#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
功能完善的点云背景平面去除脚本 (v2.1 - 终极可视化版)
作者: [你的名字或AI助手名称]
版本: 2.1
------------------------------------------------------------------------------------
本脚本提供了一个专业的工具，用于从点云数据中移除最大的背景平面。
此版本引入“爆炸视图”和强化的颜色控制，以提供无与伦比的可视化清晰度。

核心功能:
1.  自动/交互式平面分割。
2.  终极可视化:
    - “爆炸视图”功能，沿法线平移背景平面，创造物理间隙。
    - 强制颜色覆盖，确保前景/背景清晰可辨。
    - 支持对前景物体进行多种着色（统一色、原始色、高度、法线）。
    - 完全控制渲染选项（点大小、背景、坐标轴、EDL光照）。
    - 保存/加载相机视角，保存高质量截图。
3.  灵活的输出和反向选择。

依赖库:
- open3d
- numpy
- matplotlib
"""

import open3d as o3d
import numpy as np
import argparse
import time
import os
import json
from matplotlib import cm

# --- Helper Functions (保持不变) ---
def color_pcd_by_axis(pcd, axis='z', cmap_name='viridis'):
    points = np.asarray(pcd.points)
    if axis == 'x': values = points[:, 0]
    elif axis == 'y': values = points[:, 1]
    else: values = points[:, 2]
    if (np.max(values) - np.min(values)) > 1e-6:
      normalized_values = (values - np.min(values)) / (np.max(values) - np.min(values))
    else:
      normalized_values = np.zeros_like(values)
    cmap = cm.get_cmap(cmap_name)
    colors = cmap(normalized_values)[:, :3]
    pcd.colors = o3d.utility.Vector3dVector(colors)
    return pcd

# --- Core Segmentation Logic (保持不变) ---
def run_automatic_segmentation(pcd, distance_threshold, ransac_n, num_iterations):
    print("INFO: 正在以自动模式运行平面分割 (RANSAC)...")
    plane_model, inliers = pcd.segment_plane(distance_threshold=distance_threshold, ransac_n=ransac_n, num_iterations=num_iterations)
    if not inliers: raise RuntimeError("自动模式未能找到任何平面。")
    [a, b, c, d] = plane_model
    print(f"INFO: 找到平面模型: {a:.2f}x + {b:.2f}y + {c:.2f}z + {d:.2f} = 0")
    return plane_model, inliers

def run_interactive_segmentation(pcd, distance_threshold):
    print("------------------------- 交互模式说明 -------------------------")
    print("1. 在弹出的窗口中，使用 [Shift + 鼠标左键] 选取至少3个在目标平面上的点。")
    print("2. 选好后，按 [Q] 关闭窗口以继续处理。")
    print("-----------------------------------------------------------------")
    vis = o3d.visualization.VisualizerWithEditing()
    vis.create_window(window_name="交互式平面点选"); vis.add_geometry(pcd); vis.run(); vis.destroy_window()
    picked_indices = vis.get_picked_points()
    if len(picked_indices) < 3: raise RuntimeError(f"交互模式需要至少选择3个点，但只选择了 {len(picked_indices)} 个。")
    print(f"INFO: 已选择 {len(picked_indices)} 个点。正在基于这些点拟合平面...")
    picked_points = pcd.select_by_index(picked_indices)
    plane_model, _ = picked_points.segment_plane(distance_threshold=0.01, ransac_n=3, num_iterations=1000)
    points_np = np.asarray(pcd.points)
    a, b, c, d = plane_model
    distances = np.abs(a * points_np[:, 0] + b * points_np[:, 1] + c * points_np[:, 2] + d)
    inliers = np.where(distances < distance_threshold)[0].tolist()
    if not inliers: raise RuntimeError("交互模式未能根据所选点找到平面内点。")
    print(f"INFO: 找到平面模型: {a:.2f}x + {b:.2f}y + {c:.2f}z + {d:.2f} = 0")
    return plane_model, inliers


# --- 全新升级的可视化函数 ---
def ultimate_visualize_results(foreground, background, plane_model, args):
    """
    终极版可视化函数，支持“爆炸视图”和强制颜色覆盖。
    """
    print("INFO: 正在准备终极版可视化窗口...")
    
    # 1. 准备几何体
    geometries = []

    # --- 修复核心：强制颜色覆盖 ---
    # 对前景物体应用指定的着色方案，覆盖原有颜色
    if args.vis_color == 'uniform':
        foreground.paint_uniform_color([0.1, 0.9, 0.1]) # 鲜艳的绿色
    elif args.vis_color == 'z_axis':
        foreground = color_pcd_by_axis(foreground, 'z')
    elif args.vis_color == 'normal':
        if not foreground.has_normals():
            print("INFO: 前景点云没有法线，正在估算...")
            foreground.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
        foreground.normalize_normals()
        colors = (np.asarray(foreground.normals) + 1.0) / 2.0
        foreground.colors = o3d.utility.Vector3dVector(colors)
    # 如果是 'original', 则不作任何处理，保留原始颜色

    # 背景平面始终强制为灰色
    background.paint_uniform_color([0.7, 0.7, 0.7])
    
    # --- 修复核心：“爆炸视图” ---
    if args.explode_factor > 0:
        # 提取平面法向量
        a, b, c, _ = plane_model
        normal = np.array([a, b, c])
        # 计算平移向量
        translation_vector = normal * args.explode_factor
        # 平移背景点云
        background.translate(translation_vector)
        print(f"INFO: 已启用爆炸视图，背景平面沿法线平移了 {args.explode_factor} 个单位。")

    geometries.append(foreground)
    geometries.append(background)

    if args.show_axes:
        geometries.append(o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=[0, 0, 0]))

    # --- 渲染流程 (与之前版本相同) ---
    vis = o3d.visualization.Visualizer()
    vis.create_window("终极分割结果可视化", width=1280, height=720)
    opt = vis.get_render_option()
    bg_colors = {'dark': [0.1, 0.1, 0.1], 'light': [1.0, 1.0, 1.0], 'gray': [0.8, 0.8, 0.8]}
    opt.background_color = np.asarray(bg_colors[args.bg_color])
    opt.point_size = args.point_size
    if args.edl: opt.point_color_option = o3d.visualization.PointColorOption.EyeDomeLighting
    for geom in geometries: vis.add_geometry(geom)
    if args.load_view:
        try:
            ctr = vis.get_view_control()
            params = o3d.io.read_pinhole_camera_parameters(args.load_view)
            ctr.convert_from_pinhole_camera_parameters(params)
            print(f"INFO: 已从 '{args.load_view}' 加载相机视角。")
        except Exception as e: print(f"ERROR: 加载视角失败: {e}")

    print("\nINFO: 正在打开可视化窗口。按 [Q] 关闭。")
    vis.run()

    if args.save_view:
        ctr = vis.get_view_control()
        params = ctr.convert_to_pinhole_camera_parameters()
        o3d.io.write_pinhole_camera_parameters(args.save_view, params)
        print(f"INFO: 相机视角已保存到 '{args.save_view}'。")
    if args.save_image:
        vis.capture_screen_image(args.save_image, do_render=True)
        print(f"INFO: 视图已保存为图像 '{args.save_image}'。")
    vis.destroy_window()

def main():
    parser = argparse.ArgumentParser(description="一个功能完善的点云背景平面去除脚本 (v2.1)。", formatter_class=argparse.RawTextHelpFormatter)
    
    # --- 基本参数和算法参数 (保持不变) ---
    parser.add_argument("input_file", help="输入的点云文件路径。")
    parser.add_argument("output_foreground_file", help="输出的前景点云文件路径。")
    parser.add_argument("--mode", type=str, required=True, choices=['auto', 'interactive'], help="选择操作模式。")
    plane_group = parser.add_argument_group('平面分割参数')
    plane_group.add_argument('--distance_threshold', type=float, default=0.01, help="RANSAC中点到平面的最大距离。")
    # ... 其他旧参数 ...
    other_group = parser.add_argument_group('其他功能选项')
    other_group.add_argument('--output_plane_file', type=str, default=None, help="(可选) 将移除的背景平面保存到此文件。")
    other_group.add_argument('--invert', action='store_true', help="反向选择：保留平面，丢弃前景。")
    
    # --- 全新升级的可视化参数组 ---
    vis_group = parser.add_argument_group('高级可视化选项 (需要 --visualize 标志)')
    vis_group.add_argument('--visualize', action='store_true', help="处理后，启动高级可视化窗口。")
    vis_group.add_argument(
        '--vis_color', type=str, default='uniform',
        choices=['uniform', 'original', 'z_axis', 'normal'],
        help="""前景物体的着色方案:
  uniform  - (默认)强制使用鲜艳的统一颜色，提供最大对比度。
  original - 保留点云的原始颜色。
  z_axis   - 根据Z轴高度进行着色。
  normal   - 根据法向量方向着色。"""
    )
    vis_group.add_argument('--explode_factor', type=float, default=0.2, help="“爆炸视图”的平移因子。设置为0可禁用。此值为背景平面沿其法线平移的距离。")
    vis_group.add_argument('--point_size', type=float, default=2.0, help="可视化中的点大小。")
    vis_group.add_argument('--bg_color', type=str, default='dark', choices=['dark', 'light', 'gray'], help="可视化背景颜色。")
    vis_group.add_argument('--show_axes', action='store_true', help="在可视化中显示世界坐标轴。")
    vis_group.add_argument('--edl', action='store_true', help="启用Eye-Dome Lighting着色以增强深度感。")
    vis_group.add_argument('--save_image', type=str, default=None, help="将最终视图保存为图像文件。")
    vis_group.add_argument('--save_view', type=str, default=None, help="将相机视角保存到文件。")
    vis_group.add_argument('--load_view', type=str, default=None, help="从文件加载相机视角。")

    args = parser.parse_args()

    # --- 主流程 ---
    print(f"INFO: 正在从 '{args.input_file}' 加载点云...")
    try:
        pcd = o3d.io.read_point_cloud(args.input_file)
        if not pcd.has_points(): raise ValueError("点云为空。")
    except Exception as e: print(f"ERROR: 加载文件失败: {e}"); return
    
    start_time = time.time()
    try:
        if args.mode == 'auto':
            plane_model, inliers = run_automatic_segmentation(pcd, args.distance_threshold, 3, 1000)
        else:
            plane_model, inliers = run_interactive_segmentation(pcd, args.distance_threshold)
    except RuntimeError as e: print(f"ERROR: {e}"); return
    
    plane_cloud = pcd.select_by_index(inliers)
    foreground_cloud = pcd.select_by_index(inliers, invert=True)

    if args.invert: final_foreground, final_background = plane_cloud, foreground_cloud
    else: final_foreground, final_background = foreground_cloud, plane_cloud
    end_time = time.time()

    # --- 报告 & 保存 (保持不变) ---
    print("\n----------- 分割结果报告 -----------")
    print(f"处理时间      : {end_time - start_time:.4f} 秒")
    print(f"前景点数      : {len(final_foreground.points)}")
    print(f"背景点数(移除) : {len(final_background.points)}")
    print("------------------------------------")
    o3d.io.write_point_cloud(args.output_foreground_file, final_foreground)
    if args.output_plane_file: o3d.io.write_point_cloud(args.output_plane_file, final_background)

    # --- 调用全新的终极可视化函数 ---
    if args.visualize:
        ultimate_visualize_results(final_foreground, final_background, plane_model, args)
        
    print("\nINFO: 脚本执行完毕。")

if __name__ == '__main__':
    main()